#include "unifiedMathFunctions.hpp"
#include "eigenHelpers.hpp"
#include "sweSolver.hpp"
#include "reconstructionHelpers.hpp"

namespace hms::reconstruction {

template<typename SolverType, typename A, typename B>
void firstOrder(
	A& d,
	A& z,
	A& h,
	Ref<Array2Xs> v,
	const B& dCell,
	const B& zCell,
	const Ref<const Array2Xs>& vCell
){
	using S = SolverType;
	auto zeroVel = [&](){
		if constexpr ( isDenseType<remove_qualifiers<A>>() ){
			return scalar{0};
		} else {
			return Array2s{0,0};
		}
	};
	auto cellVel = [&](){
		if constexpr ( isDenseType<remove_qualifiers<A>>() ){
			return vCell;
		} else {
			return Array2s{vCell};
		}
	};

	z = zCell;
	d = select(dCell < S::minWaterDepth, 0, dCell);
	h = d + z;
	v = colwiseSelect( dCell < S::minWaterDepth, zeroVel(), cellVel() );
}

template<typename SolverType>
void firstOrder(
	EdgeValues& ev,
	const RowArray2s& dCell,
	const RowArray2s& zCell,
	const ArrayNNs<2,2>& vCell
){
	firstOrder<SolverType>(ev.dLR, ev.zLR, ev.hLR, ev.vLR, dCell, zCell, vCell);
}

template<typename SolverType, typename A, typename B>
void firstOrder(
	EdgesValues& ev,
	const Pair<A>& dCellLR,
	const Pair<A>& zCellLR,
	const Pair<B>& vCellLR
){
	static_assert( std::is_base_of_v<SolverBase<SolverType>, SolverType> );
	using S = SolverType;

	ASSERT_MATCHING_SIZES(dCellLR, zCellLR, vCellLR);

	for (int i{0}; i<2; ++i){
		reconstruction::firstOrder<S>(
			ev.dLR[i], ev.zLR[i], ev.hLR[i], ev.vLR[i],
			dCellLR[i], zCellLR[i], vCellLR[i]
		);
	}
}

template<typename SolverType, typename ES, typename EV, typename CS, typename CV>
void central(
	ES& dLR,
	ES& zLR,
	ES& hLR,
	EV& vLR,
	const CS& dCell,
	const CS& hCell,
	const CV& qCell
){
	using S = SolverType;
	using namespace reconstruction;

	/* water depth and elevation are averaged between the two adjacent cell values */
	setLR(dLR, meanLR(dCell) );
	setLR(hLR, meanLR(hCell) );
	/* bottom elevation is reconstructed from there */
	setLR(zLR, left(hLR) - left(dLR));

	/* however, for small water depths, edge values for d and v are set to zero */
	setLR( dLR, select(left(dLR) < S::minWaterDepth, 0, left(dLR) ) );
	setLR( vLR, select( // use colwiseSelect here?
		repl<2>(left(dLR) < S::minWaterDepth),
		0,
		meanLR(qCell) / repl<2>( left(dLR) )
	) );
	/* and the water elevation is set to the bottom elevation */
	setLR( hLR, select( left(dLR) < S::minWaterDepth, left(zLR), left(hLR) ) );
}

template<typename SolverType, typename EV, typename CS, typename CV>
void central(
	EV& ev,
	const CS& dCell,
	const CS& hCell,
	const CV& qCell
){
	central<SolverType>(ev.dLR, ev.zLR, ev.hLR, ev.vLR, dCell, hCell, qCell);
}



template<typename P, typename R>
void setSlopeLimiter(
	P& psi,
	const R& r,
	const LimiterFunction& limFunc
){
	using namespace reconstruction;
	using LF = LimiterFunction;

	switch(limFunc){
		case LF::minMod:
			psi = min(1, max(0, r) );
			break;
		case LF::vanAlbada:
			psi = max( 0, ( r + square(r) ) / ( 1 + square(r) ) );
			break;
		case LF::none:
			/* psi(r) is only used as a factor, so the neutral element is 1 */
			psi = 1;
			break;
		default:
			throw std::runtime_error( "Limiter function not implemented!" );
	}
}


#ifndef BLOCKWISE_INTEGRATED_PRECALC_DIFFS

/**
 * @brief Carries out the MUSCL scheme using the provided cell values, distance
 * quotients and limiter function. Results are written to the first parameter,
 * phiLR.
 * 
 * 
 * @tparam LR_w 
 * @tparam LR_r 
 * @tparam Dx 
 * @param phiLR is the out-parameter holding the results
 * @param psiLR is a buffer parameter, required to write smoothness monitor
 * and limiter function values.
 * @param phiCellLR are the cell values left and right of the edge
 * @param phiCellLLRR are the higher order neighbour cell values 
 * @param dxQuotientL (xR - xL) / (xL - xLL)
 * @param dxQuotientR (xR - xL) / (xRR - xR)
 * @param limFunc is an enum value deciding which limiter function to use
 */
template<typename LR_w, typename LR_r, typename Dx>
void musclScheme(
	LR_w& phiLR,
	LR_w& psiLR,
	const LR_r& phiCellLR,
	const LR_r& phiCellLLRR,
	const Dx& dxQuotientL,
	const Dx& dxQuotientR,
	const LimiterFunction& limFunc
)
{
	using namespace reconstruction;
	constexpr Index nRows { nRowsLR<LR_w>() };

	/* we write rLR to phiLR because the size fits */
	LR_w& rLR {phiLR};

	/* 
	 * The equations for the smoothness monitors rL and rR are:
	 *
	 *      phiL - phiLL   xR - xL         phiR - phiRR   xL - xR
	 * rL = ------------ * -------- ; rR = ------------ * --------
	 *      phiR - phiL    xL - xLL        phiL - phiR    xR - xRR
	 * 
	 * By expanding both quotients in rR with -1, we can flip them to get
	 * 
	 *      phiRR - phiR   xR - xL
	 * rR = ------------ * --------
	 *      phiR - phiL    xRR - xR
	 * 
	 * Therefore, we always subtract from right to left.
	 * Therefore, the right quotient of both rL and rR can be directly 
	 * represented by RectCells::dx() and RectCells::dy(),
	 * which already hold the distances between centroids.
	 * The right quotient is called dxQuotient below.
	 */
	/* first, we set both sides to phiR - phiL, because we have to check it
	 * for zero values later and need it on both sides */
	setLR( rLR, right(phiCellLR) - left(phiCellLR) );

	for (int i{0}; i<2; ++i){
		side(rLR, i) = select(
			side(rLR, i) == 0, 0,
			( side(phiCellLR,  i) - side(phiCellLLRR, i) )
			/ side(rLR, i)
			* repl<nRows>( i == 1 ? dxQuotientR : dxQuotientL )
			* (i == 1 ? -1 : 1) // compiler will eliminate branch, e.g. 1 - 2*i
		);
	}

	/* map to the same memory space, but as one array, to simplify access */
	auto&&
		psi { mapOrRefLR<ArrayNXsMap<nRows>>(psiLR) },
		r   { mapOrRefLR<ArrayNXsMap<nRows>>(rLR  ) };

	setSlopeLimiter( psi, r, limFunc );

	for (int i{0}; i<2; ++i){
		side(phiLR,i) =
			side(phiCellLR,i) + 0.5 * side(psiLR,i) *
			( side(phiCellLR,!i) - side(phiCellLR,i) );
	}
}

template<typename SolverType,
	typename EV,
	typename LR_r1, typename LR_r2,
	typename Dx>
void highOrder(
	EV& ev,
	const LR_r1& dCellLR,
	const LR_r1& dCellLLRR,
	const LR_r1& hCellLR,
	const LR_r1& hCellLLRR,
	const LR_r2& qCellLR,
	const LR_r2& qCellLLRR,
	const Dx& dxL,
	const Dx& dxR,
	const LimiterFunction& limFunc
){
	using S = SolverType;
	using reconstruction::musclScheme;

	using Map1 = Array1XsMap;
	using Map2 = Array2XsMap;
	using LR_w1 = decltype(ev.dLR);
	using LR_w2 = decltype(ev.vLR);
	MapOrRef<Map1, LR_w1>
		d { mapOrRefLR<Map1>(ev.dLR) },
		z { mapOrRefLR<Map1>(ev.zLR) },
		h { mapOrRefLR<Map1>(ev.hLR) };
	MapOrRef<Map2, LR_w2>
		v { mapOrRefLR<Map2>(ev.vLR) };

	/* we use the same memory space for q as for v, because the size is
	 * always the same and all operations are coeff-wise */
	auto& q {v};
	auto& qLR { ev.vLR };

	musclScheme( ev.dLR, ev.psiDH, dCellLR, dCellLLRR, dxL, dxR, limFunc );
	musclScheme( ev.hLR, ev.psiDH, hCellLR, hCellLLRR, dxL, dxR, limFunc );
	musclScheme(    qLR, ev.psiQ , qCellLR, qCellLLRR, dxL, dxR, limFunc );

	z = h - d;
	v = colwiseSelect( d < S::minWaterDepth, 0, q.rowwise() / d );
	d = select( d < S::minWaterDepth, 0, d );
	h = select( d < S::minWaterDepth, z, h );
}

#endif


template<typename EV>
void hydrostatic_firstOrder( EV& ev ){
	using namespace reconstruction;

	using M = Array1XsMap;
	using R = decltype(ev.dLR);
	MapOrRef<M,R>
		d { mapOrRefLR<M>(ev.dLR) },
		z { mapOrRefLR<M>(ev.zLR) },
		h { mapOrRefLR<M>(ev.hLR) };

	/* take the maximum value between left and right side
	 * and assign it to both sides */
	setLR( ev.zLR, max( left(ev.zLR), right(ev.zLR) ) );

	/* reconstruct water depth from water elevation minus bottom elevation, 
	 * but it mustn't be negative */
	d = max(h - z, 0);
}


} // namespace hms::reconstruction

#ifdef BLOCKWISE_INTEGRATED_PRECALC_DIFFS
#include "reconstruction_precalc_diffs.hpp"
#endif

